package top.cardone.data.jdbc.datasource.looku.impl;

import com.google.common.collect.Lists;
import lombok.Synchronized;
import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;
import org.springframework.util.CollectionUtils;
import top.cardone.data.jdbc.datasource.looku.RoutingDataSource;

import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * @author yao hai tao
 */
public class RoutingDataSourceImpl extends AbstractRoutingDataSource implements RoutingDataSource {
    private AtomicInteger counter = new AtomicInteger(1);
    /**
     * keyListThreadLocal
     */
    private ThreadLocal<List<Object>> keyListThreadLocal = new ThreadLocal<List<Object>>() {
        @Override
        protected List<Object> initialValue() {
            return Lists.newArrayList();
        }
    };
    private Object[] targetDataSourcesKeys;

    @Override
    public void addKey(Object key) {
        keyListThreadLocal.get().add(key);
    }

    @Override
    protected Object determineCurrentLookupKey() {
        if (CollectionUtils.isEmpty(keyListThreadLocal.get())) {
            return null;
        }

        //按照顺序选择从库
        int index = counter.incrementAndGet() % targetDataSourcesKeys.length;

        if (index < 0) {
            index = -index;
        }

        return targetDataSourcesKeys[index];
    }

    @Synchronized
    @Override
    public void removeKey() {
        int removeIndex = keyListThreadLocal.get().size() - 1;

        if (removeIndex < 1) {
            keyListThreadLocal.remove();

            return;
        }

        keyListThreadLocal.get().remove(removeIndex);
    }

    @Override
    public void removeKeyList() {
        keyListThreadLocal.remove();
    }

    @Override
    public void setTargetDataSources(Map<Object, Object> targetDataSources) {
        super.setTargetDataSources(targetDataSources);

        targetDataSourcesKeys = targetDataSources.keySet().toArray();
    }
}